{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Json Generation Tutorial:\n",
    "In this tutorial, we'll go over how to run the json_generation.py script, which takes a given brain parcellation nifti file and registers it to MNI space. Json files containing relevant labels and coordinates are also generated for each parcellation. This script is also able to generate json files without altering the input parcellation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First the relevant libraries must be imported. The json_generation script utilizes the nibabel, numpy, json, csv, glob, subprocess, nilearn, argpars, and matplotlib libraries.\n",
    "It also uses AFNI and FSL functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nb\n",
    "import numpy as np\n",
    "import json\n",
    "import csv\n",
    "import os\n",
    "import glob\n",
    "import subprocess\n",
    "import sys\n",
    "from datetime import datetime\n",
    "from nilearn import plotting as nip\n",
    "from nilearn import image\n",
    "from argparse import ArgumentParser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we have to specify the inputs. For the purpose of this jupyter notebook sys.argv is used to store the inputs, but when running the script in the terminal:\n",
    "\n",
    "python $<$path_to_script$>$/json_generation.py $<$input_file$>$ $<$output_dir$>$ --fig_name $<$fig_name$>$ --voxel_size $<$vox_size$>$\n",
    "\n",
    "Input Description:\n",
    "input_file = The path of the mri parcellation file you intend to process.\n",
    "        If you only specify this value, 'output_dir', and optionally '--label_csv',\n",
    "        then a JSON file will be generate without manipulating this input file.\n",
    "        \n",
    "output_dir = The directory to store the generated nii.gz and/or json file(s).\n",
    "\n",
    "output_name = Name assigned to both the processed parcellation file and its corresponding,\n",
    "        json label file. Do not include the file type (nii.gz, nii, or json).\n",
    "        If None, 'input_file' and 'reference_brain' names will be combined to\n",
    "        make $<$input_file$>$\\_$<$ref_brain$>$.nii.gz or .json. Default is None.\n",
    "        \n",
    "ref_brain = Path for reference image you wish to register your parcellation too,\n",
    "        be sure that it has the correct voxel size you want your output atlas file to have.\n",
    "        If None, resampling and registration will not be done. Default is None.\n",
    "        \n",
    "voxel_size = Voxel size (1,2,4 mm^3, etc.) of the ref_brain image specified in 'ref_brain'. This value must be\n",
    "        input if you wish to run resampling and registration. Default is 1.\n",
    "\n",
    "label_csv = csv file containing the ROI label information for the parcellation file, default is None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file = '/Users/ross/Documents/neuroparc/atlases/label/Human/AAL_space-MNI152NLin6_res-1x1x1.nii.gz'\n",
    "output_dir = '/Users/ross/Documents/neuroparc/atlases'\n",
    "output_name = 'json_tutorial'\n",
    "ref_brain = ''\n",
    "voxel_size = ''\n",
    "label_csv = ''\n",
    "\n",
    "#Necessary for running this function in a jupyter notebook\n",
    "sys.argv = ['',input_file, output_dir, '--output_name', output_name]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define the function get_centers, which takes the given atlas file and finds the center coordinates for each of the ROIs based on voxel value. So it looks at the number of voxels with a value of 1, then moves on to ones with a value of 2, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_centers(brain, orig_labs):\n",
    "    \"\"\"\n",
    "    Get coordinate centers given a nifti image loaded with nibabel\n",
    "    \n",
    "    Returns a dictionary of label: coordinate as an [x, y, z] array\n",
    "    \"\"\"\n",
    "    dat = brain.get_data()\n",
    "    \n",
    "    labs, size = np.unique(dat, return_counts=True)\n",
    "\n",
    "    size=dict(zip(labs,size))\n",
    "        \n",
    "    # Bit of a clumsy stop-gap for correcting for lost ROIs due to resampling/registration\n",
    "    for n in orig_labs:\n",
    "        if not size.get(n):\n",
    "            size[n] = None\n",
    "\n",
    "    coords_connectome = []\n",
    "    for lab in labs:\n",
    "        fd_dat = np.asarray(dat == lab).astype('float64')\n",
    "        parcel = nb.Nifti1Image(dataobj=fd_dat, header=brain.header, affine=brain.affine)\n",
    "        coords_connectome.append(nip.find_xyz_cut_coords(parcel))\n",
    "\n",
    "    return dict(zip(labs, coords_connectome)), size\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that everything has been defined, we can run the script. In this main function, if \"--ref_brain\" is defined, the input file is resampled to the deisred voxel size and registered to the defined reference anatomical nifti file. If you have a csv file containing anatomical labels for different ROI values in the parcellation input file you are using (see the files in the label/Human/Anatomical-labels-csv for examples), you can specify its location with the \"--label_csv\" flag. The information inside that csv file will be read and added to the output json metadata file for the input parcellation file.\n",
    "\n",
    "If you just want to generate the json file without resampling/registering/changing the input file in any way, do not define \"--ref_brain\". This will result in only the json file being made.\n",
    "\n",
    "Json file output examples can be found in the label/Human/Metadata-json directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "\n",
    "    parser = ArgumentParser(\n",
    "        description=\"Script to take already MNI-aligned atlas images and generate json file information.\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"input_file\",\n",
    "        help=\"\"\"The path of the mri parcellation file you intend to process.\n",
    "        If you only specify this value, 'output_dir', and optionally '--label_csv',\n",
    "        then a JSON file will be generate without manipulating this input file.\"\"\",\n",
    "        action=\"store\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"output_dir\",\n",
    "        help=\"\"\"The directory to store the generated nii.gz and/or json file(s).\"\"\",\n",
    "        action=\"store\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--output_name\",\n",
    "        help=\"\"\"Name assigned to both the processed parcellation file and its corresponding,\n",
    "        json label file. Do not include the file type (nii.gz, nii, or json).\n",
    "        If None, 'input_file' and 'reference_brain' names will be combined to\n",
    "        make <input_file>_<ref_brain>.nii.gz/.json. Default is None.\"\"\",\n",
    "        action=\"store\",\n",
    "        default=None,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--ref_brain\",\n",
    "        help=\"\"\"Path for reference image you wish to register your parcellation too,\n",
    "        be sure that it has the correct voxel size you want your output atlas file to have.\n",
    "        If None, resampling and registration will not be done. Default is None.\"\"\",\n",
    "        action=\"store\",\n",
    "        default=None,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--voxel_size\",\n",
    "        help=\"\"\"Voxel size (1,2,4 mm^3, etc.) of the ref_brain image specified in 'ref_brain'. This value must be\n",
    "        input if you wish to run resampling and registration. Default is 1.\"\"\",\n",
    "        action=\"store\",\n",
    "        default='1',\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--label_csv\",\n",
    "        help=\"csv file containing the ROI label information for the parcellation file, default is None\",\n",
    "        action=\"store\",\n",
    "        default=None,\n",
    "    )\n",
    "\n",
    "\n",
    "    # and ... begin!\n",
    "    print(\"\\nBeginning neuroparc ...\")\n",
    "\n",
    "    # ---------------- Parse CLI arguments ---------------- #\n",
    "    result = parser.parse_args()\n",
    "    input_file = result.input_file\n",
    "    output_dir = result.output_dir\n",
    "    output_name = result.output_name\n",
    "    ref_brain = result.ref_brain\n",
    "    vox_size = result.voxel_size\n",
    "    csv_f = result.label_csv\n",
    "\n",
    "    #If you have an input file, reference file, and no output name\n",
    "    if input_file and ref_brain and not output_name: \n",
    "        inp = input_file.split(\"/\")[-1]\n",
    "        inp = inp.split(\".nii\")[0]\n",
    "        refp = ref_brain.split(\"/\")[-1]\n",
    "        refp = refp.split(\".nii\")[0]\n",
    "        output_name = f\"{inp}_{refp}\"\n",
    "\n",
    "    # Load and organize csv for use in json creation\n",
    "    if csv_f:\n",
    "        biglist=[]\n",
    "        with open(csv_f, newline='', encoding = 'utf-8-sig') as csvfile:\n",
    "            spamreader = csv.reader(csvfile, delimiter=',')\n",
    "            for row in spamreader:\n",
    "                biglist.append(row[0])\n",
    "                biglist.append(row[1])\n",
    "            csv_dict = {biglist[i]: biglist[i+1] for i in range(0, len(biglist), 2)}\n",
    "\n",
    "    orig = nb.load(input_file)\n",
    "    dat = orig.get_data()\n",
    "    orig_labs, _ = np.unique(dat, return_counts=True)\n",
    "\n",
    "    if ref_brain:\n",
    "        # align input file to the dataset grid of the reference brain \"master\"\n",
    "        cmd = f\"3dresample -input {input_file} -prefix {output_dir}/{output_name}.nii.gz -master {ref_brain}\"\n",
    "        subprocess.call(cmd, shell=True)\n",
    "        # Change datatype of resampled file to 768?\n",
    "        im = nb.load(f\"{output_dir}/{output_name}.nii.gz\")\n",
    "        newdat = im.get_data().astype(np.uint32)\n",
    "        im.header['datatype'] = 768\n",
    "        nb.save(nb.Nifti1Image(dataobj=newdat, header=im.header, affine=im.affine), filename=f\"{output_dir}/{output_name}.nii.gz\")\n",
    "\n",
    "        #Register atlas to reference brain\n",
    "        output_reg = f\"{output_dir}/reg_{output_name}.nii.gz\"\n",
    "        # register image to atlas\n",
    "        cmd = f\"flirt -in {output_dir}/{output_name}.nii.gz -out {output_reg} -ref {ref_brain} -applyisoxfm {vox_size} -interp nearestneighbour\"\n",
    "        subprocess.call(cmd, shell=True)\n",
    "        \n",
    "        # Change datatype of registered file to 768?\n",
    "        im = nb.load(f\"{output_reg}\")\n",
    "        newdat=im.get_data().astype(np.uint32)\n",
    "        im.header['datatype'] = 768\n",
    "        nb.save(nb.Nifti1Image(dataobj=newdat, header=im.header, affine=im.affine), filename=output_reg)\n",
    "\n",
    "    \n",
    "    if not ref_brain: #If you just want a json file to be made, outputname will = input_file name\n",
    "        if not output_name:\n",
    "            inp = input_file.split(\"/\")[-1]\n",
    "            inp = inp.split(\".nii\")[0]\n",
    "            output_name=inp\n",
    "\n",
    "        output_reg = input_file #Have the parcel_centers run on input file without any resampling/registering\n",
    "\n",
    "\n",
    "\n",
    "    jsout = f\"{output_dir}/{output_name}.json\"\n",
    "    js_contents={\"MetaData\":{},\"rois\":{}}\n",
    "    roi_sum=0\n",
    "    count=0\n",
    "    now = datetime.now()\n",
    "    parcel_im = nb.load(output_reg)\n",
    "    parcel_centers, size= get_centers(parcel_im,orig_labs)\n",
    "    if csv_f:\n",
    "    # find a corresponding json file\n",
    "        js_contents['rois'][str(0)] = {\"label\": \"empty\", \"center\":None}\n",
    "        for (k, v) in csv_dict.items():\n",
    "            k=int(k)\n",
    "            try:\n",
    "                js_contents['rois'][str(k)] = {\"label\": v, \"center\": parcel_centers[k], \"size\":int(size[k])}\n",
    "                roi_sum=roi_sum+size[k]\n",
    "                count=count+1\n",
    "            except KeyError:\n",
    "                js_contents['rois'][str(k)] = {\"label\": v, \"center\": None, \"size\": None}\n",
    "            except TypeError:\n",
    "                js_contents['rois'][str(k)] = {\"label\": v, \"center\": None, \"size\": None}\n",
    "\n",
    "            \n",
    "    else:\n",
    "        js_contents['rois'][str(0)] = {\"label\": \"empty\", \"center\":None}\n",
    "        for (k, v) in parcel_centers.items():\n",
    "            k=int(k)\n",
    "            try:\n",
    "                js_contents['rois'][str(k)] = {\"label\": None,\"center\": parcel_centers[k],\"size\":int(size[k])}\n",
    "                roi_sum=roi_sum+size[k]\n",
    "                count=count+1\n",
    "            except KeyError:\n",
    "                js_contents['rois'][str(k)] = {\"label\": None, \"center\": None, \"size\": None}\n",
    "            except TypeError:\n",
    "                js_contents['rois'][str(k)] = {\"label\": None, \"center\": None, \"size\": None}\n",
    "        \n",
    "    #Atlas-wide Metadata\n",
    "    js_contents[\"MetaData\"] = {\"AtlasName\": \"\", \"Description\": '',\n",
    "    \"Native Coordinate Space\": '', \"Hierarchical\": '', \"Symmetrical\": '',\n",
    "    \"Number of Regions\": str(count), \"Average Volume Per Region\": str(float(roi_sum/count)), \"Year Generated\": str(now.year),\n",
    "    \"Generation Method\":'', \"Source\":''}\n",
    "                \n",
    "    with open(jsout, 'w') as jso:\n",
    "        json.dump(js_contents, jso, indent=4)\n",
    "            \n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Outputs:\n",
    "After running this script to completion, you should have the following outputs (depending on settings):\n",
    "\n",
    "1. A file named Dice\\_{atlas1}\\_&\\_{atlas2}\\_res-{vox}x{vox}x{vox}.png, containing a heatmap of the dice coefficients for each ROI.\n",
    "2. A file named Dice\\_{atlas1}\\_&\\_{atlas2}\\_res-{vox}x{vox}x{vox}.csv, containing the dice coefficient matrix information for the acccompanying png file.\n",
    "\n",
    "## Common Errors:\n",
    "- Issues may arise if atlases being compared have different voxel sizes, as the overlap measured by the dice coefficient may not be accurate\n",
    "- If the atlases you are using do not contain 'space-MNI152NLin6_res-' or end in either '.nii' or '.nii.gz', issues will arise in the naming structure of the output files. Either rename your atlases or edit the first 5 lines of code of the third code cell."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.3 64-bit",
   "language": "python",
   "name": "python37364bit7aebd514ba8143f191d4fea51945eb11"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}